import time
from pathlib import Path
import cv2
import torch

# Import utilities (make sure these are correctly imported from your utils package)
from utils.utils import (
    time_synchronized, select_device, increment_path,
    scale_coords, xyxy2xywh, non_max_suppression, split_for_trace_model,
    driving_area_mask, lane_line_mask, plot_one_box, show_seg_result,
    AverageMeter,
    LoadImages
)

# Hardcoded options
weights = "data/weights/yolopv2.pt" # Path to the model weights
source = 'C:/Users/永生理想/Desktop/1.mp4'          # Input source: image, video file, or directory
img_size = 640                       # Inference image size
conf_thres = 0.3                     # Confidence threshold for detections
iou_thres = 0.45                     # IOU threshold for non-max suppression
device = '0'                         # Device to run on: 'cpu' or CUDA device ID, e.g., '0'
save_conf = False                    # Save confidences in labels
save_txt = False                     # Save results to text files
nosave = False                       # Do not save images/videos
classes = None                       # Filter by class (e.g., [0, 1, 2])
agnostic_nms = False                 # Class-agnostic NMS
project = 'runs/detect'              # Directory to save results
name = 'exp'                         # Name of the experiment
exist_ok = False                     # Overwrite existing experiment

def detect():
    # Directories and settings
    save_img = not nosave and not source.endswith('.txt')  # Save inference images/videos

    save_dir = Path(increment_path(Path(project) / name, exist_ok=exist_ok))  # Increment run
    (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # Make directory

    inf_time = AverageMeter()
    waste_time = AverageMeter()
    nms_time = AverageMeter()

    # Load model
    stride = 32
    model = torch.jit.load(weights)
    selected_device = select_device(device)
    half = selected_device.type != 'cpu'  # Half precision only supported on CUDA
    model = model.to(selected_device)

    if half:
        model.half()  # Convert to FP16
    model.eval()

    # Set data loader
    vid_path, vid_writer = None, None
    dataset = LoadImages(source, img_size=img_size, stride=stride)

    # Run inference
    if selected_device.type != 'cpu':
        model(torch.zeros(1, 3, img_size, img_size).to(selected_device).type_as(next(model.parameters())))  # Run once
    t0 = time.time()
    for path, img, im0s, vid_cap in dataset:
        img = torch.from_numpy(img).to(selected_device)
        img = img.half() if half else img.float()  # Convert to FP16/32
        img /= 255.0  # Normalize to 0.0 - 1.0

        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # Inference
        t1 = time_synchronized()
        [pred, anchor_grid], seg, ll = model(img)
        t2 = time_synchronized()

        # Waste time: the incompatibility of torch.jit.trace causes extra time consumption in demo version
        # but this problem will not appear in the official version
        tw1 = time_synchronized()
        pred = split_for_trace_model(pred, anchor_grid)
        tw2 = time_synchronized()

        # Apply NMS
        t3 = time_synchronized()
        pred = non_max_suppression(pred, conf_thres, iou_thres, classes=classes, agnostic=agnostic_nms)
        t4 = time_synchronized()

        da_seg_mask = driving_area_mask(seg)
        ll_seg_mask = lane_line_mask(ll)

        # Process detections
        for i, det in enumerate(pred):  # Detections per image
            p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)

            p = Path(p)  # To Path
            save_path = str(save_dir / p.name)  # img.jpg
            txt_path = str(save_dir / 'labels' / p.stem) + ('')  # img.txt
            s += '%gx%g ' % img.shape[2:]  # Print image size
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # Normalization gain whwh
            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                # Write results
                for *xyxy, conf, cls in reversed(det):
                    if save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # Normalized xywh
                        if save_conf:
                            line = (cls, *xywh, conf)
                        else:
                            line = (cls, *xywh)
                        with open(txt_path + '.txt', 'a') as f:
                            f.write(('%g ' * len(line)).rstrip() % line + '\n')

                    if save_img:  # Add bounding box to image
                        plot_one_box(xyxy, im0, line_thickness=3)

            # Print time (inference)
            print(f'{s}Done. ({t2 - t1:.3f}s)')
            show_seg_result(im0, (da_seg_mask, ll_seg_mask), is_demo=True)

            # Save results (image with detections)
            if save_img:
                if dataset.mode == 'image':
                    cv2.imwrite(save_path, im0)
                    print(f" The image with the result is saved in: {save_path}")
                else:  # Video or stream
                    if vid_path != save_path:  # New video
                        vid_path = save_path
                        if isinstance(vid_writer, cv2.VideoWriter):
                            vid_writer.release()  # Release previous video writer
                        if vid_cap:  # Video file
                            fps = vid_cap.get(cv2.CAP_PROP_FPS)
                            w, h = im0.shape[1], im0.shape[0]
                        else:  # Webcam or stream
                            fps, w, h = 30, im0.shape[1], im0.shape[0]
                            save_path += '.mp4'
                        vid_writer = cv2.VideoWriter(
                            save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)
                        )
                    vid_writer.write(im0)

        # Update time meters
        inf_time.update(t2 - t1, img.size(0))
        nms_time.update(t4 - t3, img.size(0))
        waste_time.update(tw2 - tw1, img.size(0))

    print('Inference time: (%.4fs/frame)   NMS time: (%.4fs/frame)' % (inf_time.avg, nms_time.avg))
    print(f'Done. ({time.time() - t0:.3f}s)')

if __name__ == '__main__':
    with torch.no_grad():
        detect()
